package com.hzw.saas.common.storage.rsync;

import java.io.*;
import java.util.*;

import com.hzw.saas.common.storage.util.ChecksumUtil;
import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.date.SystemClock;
import cn.hutool.core.exceptions.ExceptionUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.lang.Assert;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

/**
 * 差异性校验工具
 * 匹配块：新文件中可以在旧文件中找到相匹配的文件快
 * 差异块：新文件中，两个相邻匹配块之间或文件起始与相邻匹配块之间或末尾与相邻匹配块之间的部分
 *
 * @author zzl
 * @since 06/22/2021
 */
@Slf4j
public class DiffBlockChecker {

    /**
     * 分块长度，byte
     * 需和旧文件分块大小一致
     */
    private final int blockLength;

    /**
     * 本地每个差异块分块大小
     * 小于等于0则不拆分每个差异块
     */
    private final int localBlockSplitLength;

    /**
     * 当前滚动块缓存
     */
    private final byte[] buffer;

    /**
     * 当前滚动块偏移量
     * 即滚动块第一个字节在文件中的位置
     */
    private long rollingOffset;

    /**
     * 前一滚动块偏移量
     * 主要用来判断当前滚动块是否可以通过rolling hash获取adler-32
     */
    private long preRollingOffset;

    /**
     * 当前差异块的起始偏移量
     */
    private long diffOffset;

    /**
     * 本地差异块索引
     */
    private int localBlockIndex;

    /**
     * 前一滚动块的首字节
     */
    private byte preByte;

    /**
     * 前一滚动块的adler32
     */
    private long preAdler32;

    /**
     * 旧文件分块校验和hashTable, key: weakChecksum, value: weakChecksum对应的BlockChecksum列表（key可能碰撞）
     */
    private Map<Long, List<OriginBlock>> blockChecksumMap;

    /**
     * 差异块和匹配块列表，按照在新文件中的位置排列
     */
    @Getter
    private final List<DiffBlock> diffCheckItems = new ArrayList<>();

    /**
     * 差异块列表，主要用于快速获取差异块
     */
    @Getter
    private final List<DiffBlock> localDiffCheckItems = new ArrayList<>();

    /**
     * 新文件
     */
    private final File tarFile;

    private InputStream ins;

    private RandomAccessFile raf;

    public DiffBlockChecker(File tarFile, List<OriginBlock> originBlockChecksums, int blockLength) {
        this(tarFile, originBlockChecksums, blockLength, 0);
    }

    public DiffBlockChecker(File tarFile, List<OriginBlock> originBlockChecksums, int blockLength, int localBlockSplitLength) {
        Assert.notNull(tarFile, "对比文件为null");
        Assert.isTrue(blockLength > 0, "blockLength必须大于0");
        this.tarFile = tarFile;
        this.blockLength = blockLength;
        this.localBlockSplitLength = localBlockSplitLength;
        this.buffer = new byte[this.blockLength];
        this.setBlockChecksumMap(originBlockChecksums);
    }

    public void check() {
        try {
            this.initCheckParams();
            this.initFileStream();
            this.checkDiff();
        } catch (Exception e) {
            throw ExceptionUtil.wrapRuntime(e);
        } finally {
            this.close();
        }
    }

    private void checkDiff() throws IOException {
        log.debug("check diff start...");
        long start = SystemClock.now();
        while (this.rollingOffset < this.tarFile.length()) {
            // 1. 读取当前block缓存
            log.debug("------------------{blockOffset: {}}---------------------", this.rollingOffset);
            log.debug("读取文件块缓存...");
            int readLength = RsyncIOUtil.readBuffer(raf, rollingOffset, buffer);
            if (readLength < 0) {
                log.debug("文件读取结束...");
                break;
            }
            log.debug("文件块实际长度：{}", readLength);
            // 2. 校验当前block是否匹配
            DiffBlock blockMatched = this.compareBufferChecksum(readLength);
            // 3. 若没有找到匹配
            if (Objects.isNull(blockMatched)) {
                // 3.1. 更新相关指针
                log.debug("未找到匹配块.");
                this.preRollingOffset = this.rollingOffset++;
                continue;
            }
            // 4. 若找到匹配
            // 4.1. 保存差异块信息, 若rollingOffset>diffOffset, 说明匹配块之前有差异块，保存差异块信息到diffCheckItems和localDiffCheckItems
            log.debug("找到匹配块：{index: {}}", blockMatched.getBlockId());
            if (this.rollingOffset > this.diffOffset) {
                this.saveLocalDiff(this.diffOffset, this.rollingOffset);
            }
            // 4.2. 保存匹配块信息
            log.debug("保存匹配块信息...");
            this.addDiff2ResultList(blockMatched);
            // 4.3. 更新相关指针
            this.preRollingOffset = this.rollingOffset;
            this.rollingOffset += readLength;
            this.diffOffset = this.rollingOffset;
        }
        // 5. 若diffOffset < tarFile.length, 处理末尾差异块的保存
        if (this.diffOffset < tarFile.length()) {
            this.saveLocalDiff(this.diffOffset, this.tarFile.length());
        }
        long cost = SystemClock.now() - start;
        log.debug("check diff over, cost: {}", cost);
    }

    /**
     * 保存差异性分块
     *
     * @param startOffset 起始偏移量（含）
     * @param endOffset   结束偏移量（不含）
     */
    private void saveLocalDiff(long startOffset, long endOffset) {
        // 若localBlockLength <=0，差异块不分片，直接保存
        log.debug("保存差异性分块信息：{startOffset: {}, endOffset: {}}", startOffset, endOffset);
        if (this.localBlockSplitLength <= 0) {
            DiffBlock diffCheckItem = new DiffBlock(startOffset, (int) (endOffset - startOffset), 0, String.valueOf(this.localBlockIndex++));
            this.addDiff2ResultList(diffCheckItem);
        }
        // 否则对差异块进行分片保存（此处只是保存信息）
        else {
            for (long i = startOffset; i < endOffset; i += this.localBlockSplitLength) {
                long length = i + this.localBlockSplitLength > endOffset ?
                    endOffset - i : this.localBlockSplitLength;
                DiffBlock diffCheckItem = new DiffBlock(i, (int) length, 0, String.valueOf(this.localBlockIndex++));
                this.addDiff2ResultList(diffCheckItem);
            }
        }
    }

    private void setBlockChecksumMap(List<OriginBlock> blockChecksums) {
        if (CollectionUtil.isEmpty(blockChecksums)) {
            this.blockChecksumMap = new HashMap<>();
            return;
        }
        this.blockChecksumMap = new HashMap<>(blockChecksums.size());
        for (OriginBlock blockChecksum : blockChecksums) {
            Long key = blockChecksum.getWeakChecksum();
            List<OriginBlock> value = blockChecksumMap.get(key);
            if (Objects.isNull(value)) {
                value = new ArrayList<>();
                this.blockChecksumMap.put(key, value);
            }
            value.add(blockChecksum);
        }
    }

    private void initCheckParams() {
        this.rollingOffset = 0L;
        this.preRollingOffset = 0L;
        this.diffOffset = 0L;
        this.localBlockIndex = 1;
        this.preByte = 0;
        this.preAdler32 = 0L;
        this.diffCheckItems.clear();
        this.localDiffCheckItems.clear();
    }

    private void initFileStream() throws FileNotFoundException {
        this.ins = FileUtil.getInputStream(tarFile);
        this.raf = new RandomAccessFile(tarFile, "r");
    }

    private DiffBlock compareBufferChecksum(int blockReadLength) {
        // 1. 判断前一滚动块到当前滚动块是否偏移1个字节，是则使用rolling hash计算弱校验和，否则计算当前buffer的弱校验和
        long adler32 = 0;
        if (this.rollingOffset - this.preRollingOffset == 1) {
            log.debug("计算rolling hash...");
            adler32 = ChecksumUtil.rollingHash(preAdler32, preByte, buffer[buffer.length - 1], blockReadLength);
        } else {
            log.debug("计算adler32...");
            adler32 = ChecksumUtil.adler32(buffer, 0, blockReadLength);
        }
        log.debug("分块弱校验和(adler-32): {}", adler32);
        // 记录当前块的信息
        this.preByte = buffer[0];
        this.preAdler32 = adler32;
        // 2. 判断弱校验和存在性
        List<OriginBlock> blockChecksums = blockChecksumMap.get(adler32);
        if (CollectionUtil.isEmpty(blockChecksums)) {
            log.debug("弱校验匹配失败...");
            return null;
        }
        // 3. 判断强校验和存在性
        log.debug("计算MD5...");
        String md5 = ChecksumUtil.md5Hex(buffer, 0, blockReadLength);
        log.debug("分块强校验和(MD5): {}", md5);
        OriginBlock blockMatched = blockChecksums.stream()
            .filter(blockChecksum -> Objects.equals(blockChecksum.getStrongChecksum(), md5)).findAny().orElse(null);
        log.debug("强校验和匹配{}", Objects.isNull(blockMatched) ? "失败..." : "成功.");
        return this.convert(blockMatched);
    }

    private DiffBlock convert(OriginBlock blockChecksum) {
        if (Objects.isNull(blockChecksum))
            return null;
        return new DiffBlock(blockChecksum.getBlockPos(), blockChecksum.getBlockLength(), 1, "");
    }

    private void addDiff2ResultList(DiffBlock diffCheckItem) {
        this.diffCheckItems.add(diffCheckItem);
        if (Objects.equals(diffCheckItem.getIsOrigin(), 0))
            this.localDiffCheckItems.add(diffCheckItem);
    }

    /**
     * 关闭文件流
     */
    private void close() {
        log.debug("差异检查结束，关闭文件流...");
        IoUtil.close(ins);
        IoUtil.close(raf);
    }
}
